

package com.tonyodev.fetch2okhttp;

import com.tonyodev.fetch2core.Downloader;
import com.tonyodev.fetch2core.FetchCoreUtils;
import com.tonyodev.fetch2core.InterruptMonitor;

import java.io.IOException;
import java.io.InputStream;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.ResponseBody;

import static com.tonyodev.fetch2core.FetchCoreDefaults.DEFAULT_BUFFER_SIZE;
import static com.tonyodev.fetch2core.FetchCoreUtils.getFileMd5String;
import static com.tonyodev.fetch2okhttp.OkHttpUtils.getDefaultCookieJar;

public class OkHttpDownloader implements Downloader<OkHttpClient,Request>{

    private final Map<Downloader.Response, okhttp3.Response> connections = Collections.synchronizedMap(new HashMap<Downloader.Response, okhttp3.Response>());

    private volatile OkHttpClient okHttpClient;
    private final FileDownloaderType fileDownloaderType;

    public OkHttpDownloader() {
        this(null, null);
    }

    public OkHttpDownloader(OkHttpClient okHttpClient) {
        this(okHttpClient, Downloader.FileDownloaderType.SEQUENTIAL);
    }

    public OkHttpDownloader(Downloader.FileDownloaderType fileDownloaderType) {
        this(null, fileDownloaderType);
    }

    public OkHttpDownloader(OkHttpClient okHttpClient, Downloader.FileDownloaderType fileDownloaderType) {
        this.okHttpClient = okHttpClient == null ? getOkHttpClient() : okHttpClient;
        this.fileDownloaderType = fileDownloaderType;
    }


    public OkHttpClient getOkHttpClient() {
        if (okHttpClient == null) {
            return new OkHttpClient.Builder()
                .readTimeout(20_000L, TimeUnit.MILLISECONDS)
                .connectTimeout(15_000L, TimeUnit.MILLISECONDS)
                .cache(null)
                .followRedirects(true)
                .followSslRedirects(true)
                .retryOnConnectionFailure(false)
                .cookieJar(getDefaultCookieJar())
                .build();

        }
        return okHttpClient;
    }

    public void setOkHttpClient(OkHttpClient okHttpClient) {
        this.okHttpClient = okHttpClient;
    }

    public Request onPreClientExecute(OkHttpClient client, ServerRequest request) {
        Request.Builder okHttpRequestBuilder = (new Request.Builder()).url(request.getUrl()).method(request.getRequestMethod(), null);
        for(String string :request.getHeaders().keySet()){
            okHttpRequestBuilder.addHeader(string,request.getHeaders().get(string));
        }
        return okHttpRequestBuilder.build();
    }

    private Downloader.ServerRequest getRedirectedServerRequest(Downloader.ServerRequest oldRequest, String redirectUrl) {
        return new ServerRequest(oldRequest.getId(), oldRequest.getUrl(), oldRequest.getHeaders(),
            oldRequest.getFile(), oldRequest.getFileUri(), oldRequest.getTag(),
            oldRequest.getIdentifier(), oldRequest.getRequestMethod(), oldRequest.getExtras(),
            true, redirectUrl, oldRequest.getSegment());
    }


    public Downloader.Response execute(Downloader.ServerRequest request, InterruptMonitor interruptMonitor) throws IOException {
        Request okHttpRequest = onPreClientExecute(getOkHttpClient(), request);
        if (okHttpRequest.header("Referer") == null) {
            String referer = FetchCoreUtils.getRefererFromUrl(request.getUrl());
            okHttpRequest.newBuilder().addHeader("Referer", referer).build();
        }

        okhttp3.Response okHttpResponse = okHttpClient.newCall(okHttpRequest).execute();
        Map<String, List<String>> responseHeaders = okHttpResponse.headers().toMultimap();
        int code = okHttpResponse.code();
        if (code == 302 || code == 301 || code == 303) {
            if (FetchCoreUtils.getHeaderValue(responseHeaders, new String[]{"Location"}) != null) {
                String redirectUrl = FetchCoreUtils.getHeaderValue(responseHeaders, new String[]{"Location"});
                okHttpRequest = this.onPreClientExecute(okHttpClient, getRedirectedServerRequest(request, redirectUrl != null ? redirectUrl : ""));
                if (okHttpRequest.header("Referer") == null) {
                    String referer = FetchCoreUtils.getRefererFromUrl(request.getUrl());
                    okHttpRequest.newBuilder().addHeader("Referer", referer).build();
                }
                try {
                    okHttpResponse.close();
                } catch (Exception var15) {
                }

                okHttpResponse = okHttpClient.newCall(okHttpRequest).execute();
                responseHeaders = okHttpResponse.headers().toMultimap();
                code = okHttpResponse.code();
            }
        }

        boolean success = okHttpResponse.isSuccessful();
        long contentLength = FetchCoreUtils.getContentLengthFromHeader(responseHeaders, -1L);
        ResponseBody responseBody = okHttpResponse.body();
        InputStream byteStream = responseBody != null ? responseBody.byteStream() : null;
        String errorResponseString = !success ? FetchCoreUtils.copyStreamToString(byteStream, false) : null;
        String hash = getContentHash(responseHeaders);
        boolean acceptsRanges = FetchCoreUtils.acceptRanges(code, responseHeaders);
        onServerResponse(request,
            new Downloader.Response(code, success, contentLength, null,
                request, hash, responseHeaders, acceptsRanges,
                errorResponseString));
        Downloader.Response response = new Downloader.Response(code, success, contentLength, byteStream, request, hash, responseHeaders, acceptsRanges, errorResponseString);
        connections.put(response, okHttpResponse);
        return response;
    }

    public String getContentHash(Map responseHeaders) {

        return FetchCoreUtils.getHeaderValue(responseHeaders, new String[]{"Content-MD5"}) != null ?
            FetchCoreUtils.getHeaderValue(responseHeaders, new String[]{"Content-MD5"}) : "";
    }

    public void disconnect(Downloader.Response response) {
        if (connections.containsKey(response)) {
            okhttp3.Response response1 = connections.get(response);
            connections.remove(response);
            closeResponse(response1);
        }

    }

    @Override
    public Integer getFileSlicingCount(ServerRequest request, Long contentLength) {
        return null;
    }

    public void close() {
        for (Downloader.Response response : connections.keySet()) {
            okhttp3.Response okHttpResponse = connections.get(response);
            closeResponse(okHttpResponse);
        }
        connections.clear();
    }

    private void closeResponse(okhttp3.Response response) {
        try {
            if (response != null) {
                response.close();
            }
        } catch (Exception var3) {
        }

    }

    public Downloader.FileDownloaderType getRequestFileDownloaderType(ServerRequest request, Set<FileDownloaderType> supportedFileDownloaderTypes) {
        return this.fileDownloaderType;
    }

    public boolean verifyContentHash(Downloader.ServerRequest request, String hash) {
        if (hash.isEmpty()) {
            return true;
        }
        String fileMd5 = getFileMd5String(request.getFile());
        if (fileMd5 != null && fileMd5.equals(hash)) {
            return true;
        }
        return false;
    }

    public void onServerResponse(Downloader.ServerRequest request, Downloader.Response response) {

    }

    public boolean getHeadRequestMethodSupported(Downloader.ServerRequest request) {
        return false;
    }

    public int getRequestBufferSize(Downloader.ServerRequest request) {
        return DEFAULT_BUFFER_SIZE;
    }

    public long getRequestContentLength(Downloader.ServerRequest request) {
        return FetchCoreUtils.getRequestContentLength(request, (Downloader) this);
    }

    public Set<Downloader.FileDownloaderType> getRequestSupportedFileDownloaderTypes(Downloader.ServerRequest request) {
        return FetchCoreUtils.getRequestSupportedFileDownloaderTypes(request, (Downloader)this);
    }
}
